package io.kiki.sba.registry.server.metadata.slot.arrange;

import io.kiki.sba.registry.common.model.metaserver.nodes.DataStoreNode;
import io.kiki.sba.registry.common.model.slot.Slots;
import io.kiki.sba.registry.exception.DisposeException;
import io.kiki.sba.registry.exception.InitializeException;
import io.kiki.sba.registry.exception.SofaRegistrySlotsException;
import io.kiki.sba.registry.lifecycle.Suspendable;
import io.kiki.sba.registry.lifecycle.impl.LifecycleHelper;
import io.kiki.sba.registry.observer.Observable;
import io.kiki.sba.registry.observer.impl.AbstractLifecycleObservable;
import io.kiki.sba.registry.server.metadata.MetadataLeaderService;
import io.kiki.sba.registry.server.metadata.bootstrap.config.ServerConfig;
import io.kiki.sba.registry.server.metadata.cluster.node.NodeAdded;
import io.kiki.sba.registry.server.metadata.cluster.node.NodeRemoved;
import io.kiki.sba.registry.server.metadata.lease.dataStoreNode.DataManagerObserver;
import io.kiki.sba.registry.server.metadata.lease.dataStoreNode.DataStoreNodeManagerImpl;
import io.kiki.sba.registry.server.metadata.monitor.Metrics;
import io.kiki.sba.registry.server.metadata.monitor.SlotsMonitor;
import io.kiki.sba.registry.server.metadata.slot.SlotAssigner;
import io.kiki.sba.registry.server.metadata.slot.SlotBalancer;
import io.kiki.sba.registry.server.metadata.slot.SlotManager;
import io.kiki.sba.registry.server.metadata.slot.assigner.DefaultSlotAssigner;
import io.kiki.sba.registry.server.metadata.slot.balance.DefaultSlotBalancer;
import io.kiki.sba.registry.server.metadata.slot.util.builder.SlotsBuilder;
import io.kiki.sba.registry.server.shared.comparator.NodeComparator;
import io.kiki.sba.registry.server.shared.slot.SlotsUtils;
import io.kiki.sba.registry.server.shared.util.NodeUtils;
import io.kiki.sba.registry.util.ConcurrentUtils;
import io.kiki.sba.registry.util.JsonUtils;
import io.kiki.sba.registry.util.SystemUtils;
import io.kiki.sba.registry.util.WakeUpLoopExecuteTask;
import lombok.Getter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;


@Component
public class ScheduledSlotArranger extends AbstractLifecycleObservable implements DataManagerObserver, Suspendable {
    private static final Logger logger = LoggerFactory.getLogger(AbstractLifecycleObservable.class);
    private final DataStoreNodeManagerImpl dataStoreNodeManager;

    private final SlotManager slotManager;

    private final SlotsMonitor slotsMonitor;

    private final MetadataLeaderService metadataLeaderService;

    private final Arranger arranger = new Arranger();

    private final ServerConfig serverConfig;

    private final Lock lock = new ReentrantLock();

    @Getter
    private volatile boolean SlotsProtectionMode = true;

    @Autowired
    public ScheduledSlotArranger(DataStoreNodeManagerImpl dataStoreNodeManager, SlotManager slotManager, SlotsMonitor slotsMonitor, MetadataLeaderService metadataLeaderService, ServerConfig serverConfig) {
        this.dataStoreNodeManager = dataStoreNodeManager;
        this.slotManager = slotManager;
        this.slotsMonitor = slotsMonitor;
        this.metadataLeaderService = metadataLeaderService;
        this.serverConfig = serverConfig;
    }

    @PostConstruct
    public void postConstruct() throws Exception {
        LifecycleHelper.initializeIfPossible(this);
        LifecycleHelper.startIfPossible(this);
    }

    @PreDestroy
    public void preDestroy() throws Exception {
        LifecycleHelper.stopIfPossible(this);
        LifecycleHelper.disposeIfPossible(this);
    }

    @Override
    protected void doInitialize() throws InitializeException {
        super.doInitialize();
        dataStoreNodeManager.addObserver(this);
        Thread executor = ConcurrentUtils.createDaemonThread(getClass().getSimpleName(), arranger);
        executor.start();
    }

    @Override
    protected void doDispose() throws DisposeException {
        arranger.close();
        dataStoreNodeManager.removeObserver(this);
        super.doDispose();
    }

    @Override
    public void update(Observable source, Object message) {
        logger.warn("[update] receive from [{}], message: {}", source, message);
        if (message instanceof NodeRemoved) {
            arranger.wakeup();
        }
        if (message instanceof NodeAdded) {
            arranger.wakeup();
        }
    }

    public boolean tryLock() {
        return lock.tryLock();
    }

    public void unlock() {
        lock.unlock();
    }

    private SlotsBuilder createSlotsBuilder(Slots slots, List<String> currentDataNodeIps, int slotNum, int replicas) {
        NodeComparator nodeComparator = new NodeComparator(slots.getDataServers(), currentDataNodeIps);
        SlotsBuilder slotsBuilder = new SlotsBuilder(slots, slotNum, replicas);
        slotsBuilder.init(currentDataNodeIps);

        nodeComparator.getRemoved().forEach(slotsBuilder::removeDataServerSlots);
        return slotsBuilder;
    }

    protected boolean assignSlots(SlotsBuilder slotsBuilder, Collection<String> currentDataServers) {
        Slots slots = createSlotAssigner(slotsBuilder, currentDataServers).assign();
        return refreshSlots(slots);
    }

    protected SlotAssigner createSlotAssigner(SlotsBuilder slotsBuilder, Collection<String> currentDataServers) {
        return new DefaultSlotAssigner(slotsBuilder, currentDataServers);
    }

    protected boolean balanceSlots(SlotsBuilder slotsBuilder, Collection<String> currentDataServers) {
        Slots slots = createSlotBalancer(slotsBuilder, currentDataServers).balance();
        return refreshSlots(slots);
    }

    private boolean refreshSlots(Slots slots) {
        if (slots == null) {
            logger.info("[refreshSlots] slot-table not change");
            return false;
        }
        if (!SlotsUtils.isValidSlots(slots)) {
            throw new SofaRegistrySlotsException("slot table is not valid: \n" + JsonUtils.writeValueAsString(slots));
        }
        if (slots.getEpoch() > slotManager.getSlots().getEpoch()) {
            slotManager.refresh(slots);
            return true;
        } else {
            logger.warn("[refreshSlots] slot-table epoch not change: {}", JsonUtils.writeValueAsString(slots));
            return false;
        }
    }

    protected SlotBalancer createSlotBalancer(SlotsBuilder slotsBuilder, Collection<String> currentDataServers) {
        return new DefaultSlotBalancer(slotsBuilder, currentDataServers);
    }

    @Override
    public void suspend() {
        arranger.suspend();
    }

    @Override
    public void resume() {
        arranger.resume();
    }

    @Override
    public boolean isSuspended() {
        return arranger.isSuspended();
    }

    private boolean tryArrangeSlots(List<DataStoreNode> dataStoreNodeList) {
        if (!tryLock()) {
            logger.warn("[tryArrangeSlots] tryLock failed");
            return false;
        }
        boolean modified = false;
        boolean noAssign = false;
        try {
            List<String> dataStoreNodeIpsList = NodeUtils.transferNodeToIpList(dataStoreNodeList);
            logger.info("[tryArrangeSlots][begin]arrange slot with DataNode, size={}, {}", dataStoreNodeIpsList.size(), dataStoreNodeIpsList);
            final Slots curSlots = slotManager.getSlots();
            SlotsBuilder slotsBuilder = createSlotsBuilder(curSlots, dataStoreNodeIpsList, slotManager.getSlotNums(), slotManager.getSlotReplicaNums());

            noAssign = slotsBuilder.hasNoAssignedSlots();
            if (noAssign) {
                logger.info("[re-assign][begin] assign slots to data-server");
                modified = assignSlots(slotsBuilder, dataStoreNodeIpsList);
                logger.info("[re-assign][end] modified={}", modified);
            } else
                if (slotsMonitor.isStableSlots()) {
                    logger.info("[balance][begin] balance slots to data-server");
                    modified = balanceSlots(slotsBuilder, dataStoreNodeIpsList);
                    logger.info("[balance][end] modified={}", modified);
                } else {
                    logger.info("[tryArrangeSlots][end] no arrangement");
                }
        } finally {
            unlock();
        }
        if (modified || noAssign) {
            // for log monitor
            logger.warn("[Arranging]noAssign={},modified={}", noAssign, modified);
        }
        return modified;
    }


    public boolean arrangeSync() {
        if (metadataLeaderService.amIStableAsLeader()) {
            final int minDataNodeNum = serverConfig.getDataNodeProtectionNum();
            // the start arrange with the dataNodes snapshot
            final List<DataStoreNode> dataStoreNodeList = dataStoreNodeManager.getDataStoreNodeInfo().getList();
            if (dataStoreNodeList.isEmpty()) {
                logger.warn("[Arranger] empty data server list");
                return false;
            } else {
                if (dataStoreNodeList.size() <= minDataNodeNum) {
                    SlotsProtectionMode = true;
                    logger.warn("[ProtectionMode] dataServers={} <= {}", dataStoreNodeList.size(), minDataNodeNum);
                    return false;
                }
                SlotsProtectionMode = false;
                Metrics.SlotArrange.begin();
                try {
                    return tryArrangeSlots(dataStoreNodeList);
                } finally {
                    Metrics.SlotArrange.end();
                }
            }
        } else {
            logger.info("[arrangeSync] not stable leader for arrange, leader: [{}], is-leader: [{}], isWarmup [{}]", metadataLeaderService.getLeader(), metadataLeaderService.amILeader(), metadataLeaderService.isWarmuped());
            return false;
        }
    }

    private final class Arranger extends WakeUpLoopExecuteTask {

        private final int waitingMillis = SystemUtils.getSystemInteger("registry.slot.arrange.interval.millis", 1000);

        @Override
        public int getWaitingMillis() {
            return waitingMillis;
        }

        @Override
        public void _execute_() {
            try {
                arrangeSync();
            } catch (Throwable e) {
                logger.error("failed to arrange", e);
            }
        }
    }
}
